from traits.api import *
from traitsui.api import *
from mayavi.core.api import Engine, PipelineBase
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
import numpy as np
from numpy import array, random, linspace, pi, ravel, cos, sin, empty, mgrid, arange
# from mayavi import mlab
import networkx as nx
from tvtk.api import tvtk
from mayavi.sources.builtin_surface import BuiltinSurface
import csv
from scipy.special import sph_harm

engine = Engine()
engine.start()

# engineIR = Engine()
# engineIR.start()


class Form(HasTraits):
    # engine = Instance(Engine, args=())
    # def traits_init(self):
    #     self.engine.start()

    # polydata1 = Instance(MlabSceneModel)
    # def _polydata1_default(self):
    #     return MlabSceneModel(engine=self.engine)

    polydata1 = Instance(MlabSceneModel, args={'engine': engine})

    @observe('polydata1.activated')
    def update_polydata1(self, event=None):
        mlab = self.polydata1.mlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        # scalars = np.array([[0, 1, 0]])
        mesh = tvtk.PolyData(points=points, polys=triangles)
        mesh.point_data.scalars = scalars
        # mesh.cell_data.scalars = scalars
        surf = mlab.pipeline.surface(mesh)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    polydata = Instance(MlabSceneModel, args={'engine': engine})
    # polydata = Instance(MlabSceneModel)
    # def _polydata_default(self):
    #     return MlabSceneModel(engine=self.engine)

    @observe('polydata.activated')
    def update_polydata(self, event=None):
        mlab = self.polydata.mlab

        # The numpy array data.
        points = np.array([[0, -0.5, 0], [1.5, 0, 0], [0, 1, 0], [0, 0, 0.5], [-1, -1.5, 0.1], [0, -1, 0.5], [-1, -0.5, 0], [1, 0.8, 0]], 'f')
        triangles = np.array([[0, 1, 3], [1, 2, 3], [1, 0, 5], [2, 3, 4], [3, 0, 4], [0, 5, 4], [2, 4, 6], [2, 1, 7]])
        scalars = np.random.random(points.shape)
        # The TVTK dataset.
        mesh = tvtk.PolyData(points=points, polys=triangles)
        mesh.point_data.scalars = scalars
        mesh.point_data.scalars.name = 'scalars'
        surf = mlab.pipeline.surface(mesh, opacity=0.1)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    polydata2 = Instance(MlabSceneModel, args={'engine': engine})
    # polydata2 = Instance(MlabSceneModel)
    # def _polydata2_default(self):
    #     return MlabSceneModel(engine=self.engine)

    @observe('polydata2.activated')
    def update_polydata2(self, event=None):
        mlab = self.polydata2.mlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        surf = mlab.triangular_mesh(points[:, 0], points[:, 1], points[:, 2], triangles, scalars=scalars)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )


    polydata3 = Instance(MlabSceneModel, args={'engine': engine})
    # polydata3 = Instance(MlabSceneModel)
    # def _polydata3_default(self):
    #     return MlabSceneModel(engine=self.engine)

    @observe('polydata3.activated')
    def update_polydata3(self, event=None):
        mlab = self.polydata3.mlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([[0, .5, 1]])
        mesh = tvtk.PolyData(points=points, polys=triangles)
        mesh.cell_data.scalars = scalars
        surf = mlab.pipeline.surface(mesh)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    polydata4 = Instance(MlabSceneModel, args={'engine': engine})
    # polydata4 = Instance(MlabSceneModel)
    # def _polydata4_default(self):
    #     return MlabSceneModel(engine=self.engine)

    @observe('polydata4.activated')
    def update_polydata4(self, event=None):
        with open('temperature.axdt') as file:
            fileStrings = file.readlines()
        a = fileStrings.index('[Data]\n')
        b = fileStrings.index('[Faces]\n')
        c1 = []
        for value in fileStrings[2:b]:
            if len(tmp := (value.strip().split(','))) == 4:
                c1.append(tmp)
        vt = np.array(c1).astype(float)
        v = vt[:, :3]
        t = vt[:, 3]
        c1 = []
        for value in fileStrings[b + 1:]:
            if len(tmp := (value.strip().split(','))) == 3:
                c1.append(tmp)
        f = np.array(c1).astype(int)
        mlab = self.polydata4.mlab
        surf = mlab.triangular_mesh(v[:, 0], v[:, 1], v[:, 2], f, scalars=t)
        mlab.colorbar()
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )
        pass

    v = None
    f = None
    t = None

    button = Button('显示温度场')
    buttonIR = Button('计算红外辐射')
    buttonClear = Button('清除视景')
    file = File('temperature.axdt')

    def _button_fired(self):
        with open(self.file) as file:
            fileStrings = file.readlines()
        a = fileStrings.index('[Data]\n')
        b = fileStrings.index('[Faces]\n')
        c1 = []
        for value in fileStrings[2:b]:
            if len(tmp := (value.strip().split(','))) == 4:
                c1.append(tmp)
        vt = np.array(c1).astype(float)
        v = vt[:, :3]
        t = vt[:, 3]
        c1 = []
        for value in fileStrings[b + 1:]:
            if len(tmp := (value.strip().split(','))) == 3:
                c1.append(tmp)
        f = np.array(c1).astype(int)
        self.v = v
        self.f = f
        self.t = t
        mymlab = self.polydata4.mlab
        mymlab.clf()
        surf = mymlab.triangular_mesh(v[:, 0], v[:, 1], v[:, 2], f, scalars=t)
        mymlab.colorbar()
        mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0), )
        mymlab.draw()

    def _buttonIR_fired(self):
        v = self.v
        f = self.f
        t = self.t
        thetas = np.linspace(-40, 40, 41)
        phis = np.linspace(0, 360, 181)
        ir = np.zeros((thetas.shape[0], phis.shape[0]))
        v0 = v[f[:, 0], :]
        v1 = v[f[:, 1], :]
        v2 = v[f[:, 2], :]
        tf = (t[f[:, 0]] + t[f[:, 1]] + t[f[:, 2]]) / 3
        n = np.cross(v1 - v0, v2 - v0)
        a = np.sqrt((n * n).sum(1)) / 2
        n /= 2 * a[:, np.newaxis]
        for tt, theta in enumerate(thetas):
            costheta = np.cos(theta * np.pi / 180)
            sintheta = np.sin(theta * np.pi / 180)
            print(tt)
            for pp, phi in enumerate(phis):
                cosphi = np.cos(phi * np.pi / 180)
                sinphi = np.sin(phi * np.pi / 180)
                s = np.array([costheta * cosphi, costheta * sinphi, sintheta])
                ns = n @ s
                light = ns < -1e-6
                ir[tt, pp] = np.sum(5.67e-8 * tf[light] ** 4 / np.pi * (-ns[light]) * a[light])
                # for facet in f:

                # v0=v[facet[0]]
                # v1=v[facet[1]]
                # v2=v[facet[2]]
                # n=np.cross(v1-v0,v2-v0)
                # a=np.sqrt(np.dot(n,n))/2
                # n/=2*a
                # if (sn:=np.dot(s,n))<-1e-6:
                #     ir[tt,pp]+=5.67e-8/np.pi*sn*a
        self.ir = ir
        self.thetas = thetas
        self.phis = phis
        print(ir.max())

    def _buttonClear_fired(self):
        mymlab = self.polydata4.mlab
        mymlab.clf()
        mymlab.draw()
        pass

    buttonShowIR = Button('显示红外辐射')
    buttonShowIRSphere = Button('显示红外辐射球面')

    def _buttonShowIR_fired(self):
        mymlab = self.polydata4.mlab
        # mymlab = self.sceneIR.mymlab
        thetas, phis = np.meshgrid(self.thetas, self.phis,indexing='ij')
        costhetas = np.cos(thetas * np.pi / 180)
        sinthetas = np.sin(thetas * np.pi / 180)
        cosphis = np.cos(phis * np.pi / 180)
        sinphis = np.sin(phis * np.pi / 180)
        x = costhetas * cosphis
        y = costhetas * sinphis
        z = sinthetas
        ir=self.ir*5
        # r1 = np.abs(np.sin(theta)) + 1
        mymlab.mesh(ir * x, ir * y, ir * z,opacity=.5)
        mymlab.colorbar()
    def _buttonShowIRSphere_fired(self):
        mymlab = self.polydata4.mlab
        mymlab = self.sceneIR.mymlab
        thetas, phis = np.meshgrid(self.thetas, self.phis,indexing='ij')
        costhetas = np.cos(thetas * np.pi / 180)
        sinthetas = np.sin(thetas * np.pi / 180)
        cosphis = np.cos(phis * np.pi / 180)
        sinphis = np.sin(phis * np.pi / 180)
        x = costhetas * cosphis
        y = costhetas * sinphis
        z = sinthetas
        ir=self.ir*5
        # r1 = np.abs(np.sin(theta)) + 1
        mymlab.mesh(x, y, z, scalars=ir,opacity=.5)
        mymlab.colorbar()
        # from mayavi import mymlab
        # mymlab.mesh


    sceneIR = Instance(MlabSceneModel, args={'engine': engine})
    # sceneIR = Instance(MlabSceneModel)
    # def _sceneIR_default(self):
    #     return MlabSceneModel(engine=self.engine)

    buttonShowIRScene = Button('测试')
    def _buttonShowIRScene_fired(self):
        mymlab = self.sceneIR.mlab
        mymlab.figure(self.sceneIR.mayavi_scene)
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        # mymlab.figure()
        surf = mymlab.triangular_mesh(points[:, 0], points[:, 1], points[:, 2], triangles, scalars=scalars)
        mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0), )
        mymlab.draw()

        #
        # # mymlab = self.polydata1.mymlab
        # points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        # triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        # scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        # # scalars = np.array([[0, 1, 0]])
        # mesh = tvtk.PolyData(points=points, polys=triangles)
        # mesh.point_data.scalars = scalars
        # # mesh.cell_data.scalars = scalars
        # surf = mymlab.pipeline.surface(mesh)
        # mymlab.pipeline.surface(mymlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    def traits_init(self):
        # self.polydata=MlabSceneModel(engine=engine)
        # self.sceneIR=MlabSceneModel(engine=engine)
        pass

    view = View(
        Tabbed(
            Group(
                UItem('sceneIR', editor=SceneEditor(scene_class=MayaviScene)),
                HGroup(
                    UItem('buttonShowIRScene'),
                    spring,
                ),
                label='红外辐射'
            ),
            UItem('polydata2', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('polydata', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('polydata1', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('polydata3', editor=SceneEditor(scene_class=MayaviScene)),
            Group(
                UItem('polydata4', editor=SceneEditor(scene_class=MayaviScene)),
                HGroup(
                    UItem('button'),
                    UItem('buttonIR'),
                    UItem('buttonClear'),
                    UItem('buttonShowIR'),
                    UItem('buttonShowIRSphere'),
                    UItem('file'),
                ),
                label='温度场'
            ),
        ),
        height=800,
        width=1000,
        resizable=True  # We need this to resize with the parent widget
    )

#
# from PyQt5.QtWidgets import *
#
#
# class FormWidget(QWidget):
#     def __init__(self):
#         QWidget.__init__(self)
#         layout = QHBoxLayout(self)
#         layout.setContentsMargins(0, 0, 0, 0)
#         layout.setSpacing(0)
#         self.form = Form()
#         layout.addWidget(self.form.edit_traits(parent=self, kind='subpanel').control)
#
#
# formWidget = FormWidget()
if __name__ == '__main__':
    Form().configure_traits()

    # import os
    # import sys
    #
    # app = QApplication(sys.argv)
    # formWidget.show()
    # sys.exit(app.exec_())
#